# 导入所需库
from langchain_core.prompts import PromptTemplate
from langchain_deepseek import ChatDeepSeek
import os
from dotenv import load_dotenv, find_dotenv
from langchain_openai import ChatOpenAI
from langchain.agents import create_openai_tools_agent, AgentExecutor, tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain.memory import ConversationTokenBufferMemory
from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from interview_tools import *
import uuid
import tiktoken  # 添加 tiktoken 用于 token 计数
import asyncio
from fastapi.responses import StreamingResponse
from typing import AsyncGenerator, Generator, Any
from DatabaseTool import *

# 1、使用.env 存储api_key
load_dotenv()

os.environ["LANGCHAIN_TRACING_V2"] = "false"
os.environ["LANGCHAIN_API_KEY"] = "ls__a339aed7ea934ec999736421cdb43629"
os.environ["LANGCHAIN_PROJECT"] = "samibot"

os.environ["OPENAI_API_KEY"] = "sk-83S6JB0eSdH536lC9aD1F16482C04dC69c02Ba6571B5A0Fb"
os.environ["OPENAI_API_BASE"] = "https://ai-yyds.com/v1"

os.environ["SERPAPI_API_KEY"] = "e7b15718df583476214651b32a4e835523f3bac273d511341a1de30430b2164e"
REDIS_URL = "redis://:@127.0.0.1:16379"  # os.getenv("REDIS_URL")

app = FastAPI()

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 允许所有源
    allow_credentials=True,
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)


# 自定义 DeepSeek 类添加 token 计数方法
class CustomChatDeepSeek(ChatOpenAI):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        # 使用类变量而不是实例变量
        self._encoding = None

    @property
    def encoding(self):
        """延迟加载编码器"""
        if self._encoding is None:
            self._encoding = tiktoken.get_encoding("cl100k_base")
        return self._encoding

    def get_num_tokens_from_messages(self, messages):
        """计算消息的 token 数量（兼容 OpenAI 格式）"""
        num_tokens = 0
        for message in messages:
            num_tokens += 4  # 每条消息的开销

            # 处理不同类型的消息内容
            content = message.content
            if isinstance(content, list):
                # 处理多模态内容
                for item in content:
                    if item["type"] == "text":
                        num_tokens += len(self.encoding.encode(item["text"]))
            else:
                num_tokens += len(self.encoding.encode(str(content)))

            if hasattr(message, "name") and message.name:  # 检查是否存在name属性
                num_tokens += -1  # 名字字段有额外开销
        num_tokens += 2  # 回复的开销
        return num_tokens


class InterviewServer:

    def __init__(self, session_id : str, info : dict):
        self.chatmodel = CustomChatDeepSeek(
            api_key="sk-c40237f01d7a491d9a7e792772e9e3ca",
            base_url="https://api.deepseek.com/v1",  # 注意 /v1 路径
            model="deepseek-chat",
            temperature=0,
            streaming=True,  # 确保启用流式输出
        )

        self.MEMORY_KEY = "chat_history"
        # 注释: 2. 最后可以询问一些企业相关的问题。 这个技术点暂时不一定会被触发。
        self.SYSTEMPL = f"""
        背景: 你是一个面试官,你精通如何考察面试人员对企业需求的匹配度,并会在面试中通过不断的询问问题来达到考察的目的。你的公司名字叫上海第二工业大学
        # 现阶段情况:
        1. 在你面前的面试者姓名:{info['real_name']}
        2. 在你面前的面试者需要面试的岗位:{info['job_name']}
        3. 面试职位的岗位要求:{info['job_name']}
        4. 企业对该面试经验的要求:{info['job_name']}
        # 你的面试流程:
        第一步自我介绍。
        第二步项目经验介绍
        第三步企业关心的垂直经验问题。
        # 判断面试面试流程每一个步骤结束标准
        1. 如果是基础岗位,自我介绍尽量包含需要包含学校,所属专业,实践经验等
        2. 如果是技术岗,尽可能介绍2个项目经验。
        3. 企业关心的垂直领域的问题，尽可能回答到位。
        # 面试结束信号:
        1. 问答轮数达到30轮,也就是提问和回答算1次，有30次提问和回答,则需要终止面试。
        2. 面试的第一步，第二步，第三步完成达到面试流程结束标准后，整个面试则可以结束。
        3. 面试的流程中每一步中，如果对话超过10轮后还没有把内容聊清楚，就可以进入下一个步骤。
        # 复试面试要求:
        1. 介绍尽量简明扼要,问到企业想要的要求点,面试官的回复字数不超过100字。
        2. 最后可以询问一些企业相关的问题。
        """

        self.prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    self.SYSTEMPL,
                ),
                MessagesPlaceholder(variable_name=self.MEMORY_KEY),
                (
                    "user",
                    "{input}"
                ),
                MessagesPlaceholder(variable_name="agent_scratchpad"),
            ],
        )

        tools = [test, getnowtime]
        agent = create_openai_tools_agent(
            llm=self.chatmodel,
            tools=tools,
            prompt=self.prompt
        )
        self.sessionId = session_id
        self.memory = self.get_memory(self.sessionId)
        memory = ConversationTokenBufferMemory(
            llm=self.chatmodel,
            human_prefix="用户",
            ai_prefix="ergongda",
            memory_key=self.MEMORY_KEY,
            output_key="output",
            return_messages=True,
            max_token_limit=4000,
            chat_memory=self.memory,
        )
        self.agent_executor = AgentExecutor(
            agent=agent,
            tools=tools,
            memory=memory,
            verbose=True,
            handle_parsing_errors=True  # 添加错误处理
        )

    def get_memory(self, sessid: str):
        chat_message_history = RedisChatMessageHistory(
            url=REDIS_URL, session_id=sessid
        )
        # chat_message_history.clear()#清空历史记录
        print("chat_message_history:", chat_message_history.messages)
        store_message = chat_message_history.messages
        if len(store_message) > 10:
            prompt = ChatPromptTemplate.from_messages(
                [
                    (
                        "system",
                        self.SYSTEMPL + "\n这是一段你和用户的对话记忆，对其进行总结摘要，摘要使用第一人称‘我’，并且提取其中的用户关键信息。以如下格式返回:\n 总结摘要内容｜用户关键信息 \n 例如 面试者张三学校上海交通大学，擅长编程领域，对银行项目比较有经验，擅长多线程项目开发。｜张三,生日1999年1月1日"
                    ),
                    ("user", "{input}"),
                ]
            )
            chain = prompt | self.chatmodel
            summary = chain.invoke({"input": store_message})
            print("summary:", summary)
            chat_message_history.clear()
            chat_message_history.add_message(summary)
            print("总结后：", chat_message_history.messages)
        return chat_message_history

    async def stream_response(self, query: str) -> AsyncGenerator[str, None]:
        """流式生成响应"""
        try:
            # 使用异步迭代器流式处理响应
            async for chunk in self.agent_executor.astream({"input": query, "chat_history": self.memory.messages}):
                if "output" in chunk:
                    # 确保输出是字符串
                    output = str(chunk["output"])
                    # 流式输出每个字符（真实场景中可以按token输出）
                    for char in output:
                        yield f"data: {json.dumps(char, ensure_ascii=False)}\n\n" # yield char
                        await asyncio.sleep(0.06)  # 控制输出速度
        except Exception as e:
            error_msg = f"处理错误: {str(e)}"
            yield error_msg
    # async def stream_response(self, query: str) -> AsyncGenerator[str, None]:
    #     """流式生成响应 - 修复版本"""
    #     try:
    #         # 使用异步迭代器流式处理响应
    #         async for chunk in self.agent_executor.astream({"input": query, "chat_history": self.memory.messages}):
    #             if "output" in chunk:
    #                 # 确保输出是字符串
    #                 output = str(chunk["output"])
    #                 # 发送符合SSE格式的数据
    #                 yield f"data: {output}\n\n"
    #                 await asyncio.sleep(0.06)  # 控制输出速度
    #     except Exception as e:
    #         error_msg = f"data: 处理错误: {str(e)}\n\n"
    #         yield error_msg
    #     finally:
    #         # 发送结束标志
    #         yield "event: end\ndata: [DONE]\n\n"
    # async def stream_response(self, query: str) -> AsyncGenerator[str, None]:
    #     """流式生成响应 - 完全符合SSE规范"""
    #     try:
    #         # 发送初始心跳（保持连接）
    #         yield ": heartbeat\n\n"
    #
    #         # 流式处理响应
    #         buffer = ""
    #         async for chunk in self.agent_executor.astream({"input": query, "chat_history": self.memory.messages}):
    #             if "output" in chunk:
    #                 output = str(chunk["output"])
    #
    #                 # 将输出添加到缓冲区
    #                 buffer += output
    #
    #                 # 寻找自然断点（如句子结束）
    #                 last_punctuation = max(buffer.rfind('。'), buffer.rfind('？'), buffer.rfind('！'), buffer.rfind('.'))
    #                 if last_punctuation != -1 and len(buffer) > 10:
    #                     # 发送到断点处的内容
    #                     to_send = buffer[:last_punctuation + 1]
    #                     buffer = buffer[last_punctuation + 1:]
    #
    #                     # 确保格式符合SSE规范
    #                     yield f"data: {json.dumps(to_send, ensure_ascii=False)}\n\n"
    #
    #                 # 定期发送心跳
    #                 yield ": heartbeat\n\n"
    #                 await asyncio.sleep(0.06)
    #
    #         # 发送剩余内容
    #         if buffer:
    #             yield f"data: {json.dumps(buffer, ensure_ascii=False)}\n\n"
    #
    #     except Exception as e:
    #         error_msg = f"处理错误: {str(e)}"
    #         yield f"data: {json.dumps(error_msg, ensure_ascii=False)}\n\n"
    #     finally:
    #         # 发送结束事件
    #         yield "event: end\ndata: [DONE]\n\n"

@app.get("/chat")
async def chat(query: str, context: str):
    # 初始化数据参数
    db_tool = DatabaseTool()
    info = db_tool.get_candidate_info(context) #mobile

    master = InterviewServer(session_id=context, info=info)
    # 创建流式响应 - 添加SSE所需的headers
    response = StreamingResponse(
        master.stream_response(query),
        media_type="text/event-stream"
    )
    # response.headers["Cache-Control"] = "no-cache"
    # response.headers["Connection"] = "keep-alive"
    # response.headers["X-Accel-Buffering"] = "no"
    return response


# WebSocket 流式接口
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    try:
        while True:
            data = await websocket.receive_json()
            query = data.get("query", "")
            context = data.get("context", "default_session")

            master = InterviewServer(session_id=context)

            # 流式处理并发送响应
            async for chunk in master.stream_response(query):
                await websocket.send_text(chunk)

    except WebSocketDisconnect:
        print("客户端断开连接")
    except Exception as e:
        error_msg = f"错误: {str(e)}"
        await websocket.send_text(error_msg)


'''
测试fastapi启动后是否能正常工作
'''
@app.get("/")
def test_demo():
    return {"Hello": "World11"}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)